#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <dlfcn.h>  
#include <errno.h>

#include "elf64.h"
#include "sre_memory.h"
#include "sre_printf.h"
#include "sre_socket.h"
#include "sre_error.h"
#include "sre_version.h"

typedef void (*entryFunc)();
typedef struct tag_elf_ftable
{
	char name[1024];
	void *funcPtr;
} elf_ftable_t;

typedef struct tag_sec_table
{
	uint64_t addr;
	uint64_t start;
	uint64_t end;
} sec_table_t;

typedef struct tag_so_handle
{
	char name[64];
	void *handle;
	int load;
} so_handle_t;

typedef struct {
  uint64_t base_addr;
  uint64_t initEntryAddr;
  uint64_t initArrayAddr;
  uint64_t initArraySize;
  uint64_t entryAddr; 
  uint64_t jmpRelAddr; 
  uint64_t pltRelType; 
  uint64_t pltRelSize; 
  uint64_t symTabAddr; 
  uint64_t strTabAddr; 
  uint64_t strTabSize; 
  uint64_t RelAddr; 
  uint64_t RelSize; 
  
  uint64_t soaddr[16];
  uint32_t soindex;

  sec_table_t sectable[16];
  uint32_t secidx;
} ElfSecInfo;

uint8_t *elf_mmap(uint8_t *addr, int len)
{
	if (addr) {
		int pagesize = sysconf(_SC_PAGE_SIZE);
		uint64_t start = (uint64_t)addr;
		uint64_t end = (uint64_t)addr + len;
		start = SRE_CEIL(start, pagesize);
		end = SRE_ALIGN(end, pagesize);
		int size = end - start;
        uint8_t *mem = (uint8_t *)mmap((void *)start, size, PROT_READ|PROT_WRITE|PROT_EXEC, MAP_SHARED|MAP_ANONYMOUS, -1, 0);
		//printf("mmap addr1: %p-0x%x\n", mem, size);
		memset(mem,0,size);
		//mprotect(mem, size, PROT_READ|PROT_WRITE|PROT_EXEC);
		return MAP_FAILED == mem ? NULL : addr;
	} else {
		uint8_t *mem = (uint8_t *)mmap(NULL, len, PROT_READ|PROT_WRITE|PROT_EXEC, MAP_SHARED|MAP_ANONYMOUS, -1, 0);
		//printf("mmap addr2: %p-0x%x\n", mem, len);
		memset(mem,0,len);
		//mprotect(mem, size, PROT_READ|PROT_WRITE|PROT_EXEC);
		return MAP_FAILED == mem ? NULL : mem;
	}
}

int g_argc;
char** g_argv;

int elf_start_main(int (*m)(int, char**, char**), int argc, char** argv) 
{
  if (g_argc) {
    argc = g_argc;
    argv = g_argv;
  }
  exit(m(argc, argv, 0));
}

elf_ftable_t g_ftable[] = 
{
    {"stdin", NULL},
    {"stdout", NULL},
    {"stderr", NULL},
	{"optarg", NULL},
	{"__libc_start_main", (void *)elf_start_main},
#if 1
	{"getopt", (void *)sre_getopt},
	{"exit", (void *)sre_exit},
	{"perror", (void *)sre_perror},
	{"__assert_fail", (void *)__sre_assert_fail},
	{"__errno_location", (void *)__sre_errno_location},
	{"printf", (void *)sre_printf},
	{"printf_chk", (void *)sre_printf_chk},
	{"__printf_chk", (void *)sre_printf_chk},
	{"putchar", (void *)sre_putchar},
	{"puts", (void *)sre_puts},
	{"fgets", (void *)sre_fgets},
	{"sprintf", (void *)sre_sprintf},
	{"memset", (void *)sre_memset},
	{"memcpy", (void *)sre_memcpy},
	{"memmove", (void *)sre_memmove},
	{"memcmp", (void *)sre_memcmp},
	{"malloc", (void *)sre_malloc},
	{"free", (void *)sre_free},
	{"inet_addr", (void *)sre_inet_addr},
	{"socket", (void *)sre_socket},
	{"close", (void *)sre_close},
	{"bind", (void *)sre_bind},
	{"recvfrom", (void *)sre_recvfrom},
	{"sendto", (void *)sre_sendto},
#endif
};

so_handle_t g_sohandle[] = 
{
	{"libstdc++.so.6", NULL, 1},
	{"libpthread.so.0", NULL, 1},
	{"libgcc_s.so.1", NULL, 1},
	{"libc.so.6", NULL, 1},
	{"libm.so.6", NULL, 1},
	{"libz.so.1", NULL, 1},
	//{"libdce.so", NULL, 1},
	{"libtrace.so.3", NULL, 1},
	{"libwandio.so.1", NULL, 1},
	{"libpcap.so.1", NULL, 1},
};

elf_ftable_t g_sotable[10240];
uint32_t g_pos = 0;

uint64_t get_addr(ElfSecInfo *secinfo, uint64_t addr)
{
	for (uint32_t k = 0; k < secinfo->secidx; k++) {
		if (addr >= secinfo->sectable[k].start && addr <= secinfo->sectable[k].end) {
			return (secinfo->sectable[k].addr+addr-secinfo->sectable[k].start);
		}
	}
	return addr;
}

int load_dl(void)
{
	for (int i = 0; i < sizeof(g_sohandle)/sizeof(so_handle_t); i++) {
		if (g_sohandle[i].load) {
			g_sohandle[i].handle = dlopen(g_sohandle[i].name, RTLD_LAZY);
			if (g_sohandle[i].handle == NULL) {
				printf("dlopen %s failed\n", g_sohandle[i].name);  
				return -1;
			}
		}
	}

	return 0;
}

void close_dl(void)
{
	for (int i = 0; i < sizeof(g_sohandle)/sizeof(so_handle_t); i++) {
		if (g_sohandle[i].handle) {
			dlclose(g_sohandle[i].handle);
		}
	}
}

void sre_undefined() {
  fprintf(stderr, "undefined function is called\n");
  abort();
}

void *find_dlso(char *name)
{
    void* func = NULL;
    
	if (name[0] == 0x00) {
		return func;
	}

	//printf("[find_dlso]: ---0---%s\n",name);
    for (uint32_t i=0; i < sizeof(g_ftable)/sizeof(elf_ftable_t); i++) {
      if (strcmp(name, g_ftable[i].name) == 0) {
		func = (void *)g_ftable[i].funcPtr;
		//printf("[find_dlso]: ---1---%s,%p\n",name,func);
        return func;
      }
    }


    if (!func) {
      func = dlsym(RTLD_DEFAULT, name);
	  if (func) {
		  //printf("[find_dlso]: ---2---%s,%p\n",name,func);
		  return func;  
	  }
	}

    if (!func) {
		for(uint32_t i=0;i<sizeof(g_sohandle)/sizeof(so_handle_t);i++) {
			if (g_sohandle[i].handle) {
				func = (void *)dlsym(g_sohandle[i].handle, name);
				if (func) { 
					//printf("[find_dlso]: ---3---%s,%p\n",name,func);
					return func;  
				}
			}
		}
	}

	for (uint32_t i = 0; i < g_pos; i++) {
		if (strcmp(name, g_sotable[i].name) == 0) {
			func = (void *)g_sotable[i].funcPtr;
			if (func) { 
				//printf("[find_dlso]: ---4---%s,%p\n",name, func);
				return func;  
			}
		} 
	}

	//printf("[find_dlso]: ---5---%s,%p\n",name,func);
	return func;
}

void relocate(char *title, char* reloc_type, ElfSecInfo *secinfo) 
{
  int i, k;
  int relsz = 0;
  Elf64_Rela* rel = NULL;
  char* dstr = (char*)secinfo->strTabAddr;
  Elf64_Sym* dsym = (Elf64_Sym*)secinfo->symTabAddr;

  if (strcmp(reloc_type, "rel") == 0) {
	  rel = (Elf64_Rela*)secinfo->RelAddr;
	  relsz = secinfo->RelSize;
  } else {
	  rel = (Elf64_Rela*)secinfo->jmpRelAddr;
	  relsz = secinfo->pltRelSize;
  }

  for (i = 0; i < relsz / sizeof(Elf64_Rela); rel++, i++) {
    long* addr = (long*)((char*)rel->r_offset + secinfo->base_addr);
    int type = ELF64_R_TYPE(rel->r_info);
    Elf64_Sym* sym = dsym + ELF64_R_SYM(rel->r_info);
    char* sname = dstr + sym->st_name;
	void *val = find_dlso(sname);

	//printf("[%s]: %s %p %s %d => %p\n", title, reloc_type, (void*)addr, sname, type, val);

    switch (type) {
	case R_X86_64_RELATIVE:
       *addr = (long)(secinfo->base_addr + rel->r_addend);
       //printf("[%s]: undefined rel: %s,%s\n", title, reloc_type, sname);
	   break;
    case R_X86_64_64: {
      *addr += (long)(secinfo->base_addr + rel->r_addend);
       //printf("[%s]: undefined 64: %s,%s\n", title, reloc_type, sname);
       break;
    }
    case R_X86_64_COPY: {
     // printf("[%s]: undefined copy: %s,%s\n", title, reloc_type, sname);
      if (val) {
        *addr = (long)val;
        if (strcmp(sname, "optarg") == 0) {
        	sre_optarg = addr;
        } else if (strcmp(sname, "_ZSt4cout") == 0) {
        	memcpy(addr, val, sizeof(std::cout));
        }
      } else {
        abort();
      }
      break;
    }
    case R_X86_64_GLOB_DAT: {
     // printf("[%s]: undefined glob: %s,%s\n", title, reloc_type, sname);
     *addr = (long)val;
      break;
    }
    case R_X86_64_JMP_SLOT: {
      //printf("[%s]: undefined jmp: %s,%s\n", title, reloc_type, sname);
      if (val) {
        *addr = (long)val;
      } else {
        *addr = (long)&sre_undefined;
      }
      break;
    }
    default:
      printf("[%s]: unsupported reloc: %d\n", title, type);
	  break;
    }
  }
}

uint64_t g_so_base = 0;

int LoadAndSecInfo(char *buf, Elf64_Ehdr *ehdr, ElfSecInfo *secinfo)
{
  if (ehdr->e_type == ET_DYN) {
	  g_so_base += 0x1000000;
	  secinfo->base_addr = g_so_base;
  } 

  for (int i = 0; i < ehdr->e_phnum; i++) {
    Elf64_Phdr *ephdr = (Elf64_Phdr *)(buf+ehdr->e_phoff);
    switch (ephdr[i].p_type) {
	    case PT_LOAD: {
		  uint8_t *vaddr = (uint8_t *)(ephdr[i].p_vaddr+secinfo->base_addr);
		  uint8_t* addr = (uint8_t *)elf_mmap((uint8_t*)vaddr, ephdr[i].p_memsz);
		  if (addr == MAP_FAILED) {
			 printf("[LoadAndSecInfo]: allocate memory at specific address failed\n");
			 return -1;
		  }
		  if (secinfo->secidx < 16) {
			  secinfo->sectable[secinfo->secidx].addr = (uint64_t)addr;
			  secinfo->sectable[secinfo->secidx].start = ephdr[i].p_vaddr;
			  secinfo->sectable[secinfo->secidx++].end = ephdr[i].p_vaddr + ephdr[i].p_memsz;
		  } else {
			  printf("[LoadAndSecInfo]: secidx > 16\n");
		  }

		  printf("[LoadAndSecInfo]: allocate address at: %p, 0x%lx, 0x%lx\n", addr, ephdr[i].p_vaddr, ephdr[i].p_memsz);
		  memcpy(addr, buf+ephdr[i].p_offset, ephdr[i].p_filesz);
	      break;
	    }
	    case PT_DYNAMIC: {
		  Elf64_Dyn *dynSection = (Elf64_Dyn*)(buf+ephdr[i].p_offset);

	      for (; dynSection->d_tag != DT_NULL; dynSection++) {
	   		uint64_t dynSectionAddr = get_addr(secinfo, dynSection->d_un.d_val);

	        switch (dynSection->d_tag) {
		        case DT_NEEDED: {
				  if (secinfo->soindex < 16) {
					  secinfo->soaddr[secinfo->soindex++] = dynSection->d_un.d_val;
				  } else {
					  printf("[LoadAndSecInfo]: soindex > 16\n");
				  }
		          break;
		        }
		        case DT_INIT: {
		          secinfo->initEntryAddr = dynSectionAddr;
		          break;
		        }
		        case DT_INIT_ARRAY: {
		          secinfo->initArrayAddr = dynSectionAddr;
		          break;
		        }
		        case DT_INIT_ARRAYSZ: {
		          secinfo->initArraySize = dynSection->d_un.d_val;
		          break;
		        }
		        case DT_PLTRELSZ: {
		          secinfo->pltRelSize = dynSection->d_un.d_val;
		          break;
		        }
		        case DT_STRTAB: {
		          secinfo->strTabAddr = dynSectionAddr;
		          break;
		        }
		        case DT_STRSZ: {
		          secinfo->strTabSize = dynSection->d_un.d_val;
		          break;
		        }
		        case DT_JMPREL: {
		          secinfo->jmpRelAddr = dynSectionAddr;
		          break;
		        }
		        case DT_SYMTAB: {
		          secinfo->symTabAddr = dynSectionAddr;
		          break;
		        }
		        case DT_REL:
		        case DT_RELA: {
		          secinfo->RelAddr = dynSectionAddr;
		          break;
		        }
		        case DT_RELSZ:
		        case DT_RELASZ: {
		          secinfo->RelSize = dynSection->d_un.d_val;
		          break;
		        }
		        case DT_RELENT: {
		          int relent = dynSection->d_un.d_val;
		          break;
		        }
		        case DT_PLTREL: {
		          int pltRelType = dynSection->d_un.d_val;
		          break;
		        }
		        default: {
		         // printf("[LoadAndSecInfo]: unknown DYN %d %d\n", dynSection->d_tag, dynSection->d_un.d_val);
				  break;
				}
	        }
	      }
	      
          if (!secinfo->symTabAddr || !secinfo->strTabAddr) {
             printf("[LoadAndSecInfo]: no dsym or strTabAddr\n");
		     return -1;
	      }
	      
		  break;
	    }
	    default:
	     // printf("[LoadAndSecInfo]: unknown PT %d\n", phdr[i].p_type);
		  break;
      }
  }

  secinfo->entryAddr = get_addr(secinfo, ehdr->e_entry);

  return 0;
}

int load_sym(char *buf, Elf64_Ehdr *ehdr, ElfSecInfo *secinfo)
{
  Elf64_Shdr *shdrs = (Elf64_Shdr *)(buf + ehdr->e_shoff);
  int shnum = ehdr->e_shnum == 0 ? shdrs[0].sh_size : ehdr->e_shnum;
  int shstrndx = ehdr->e_shstrndx == SHN_XINDEX ? shdrs[0].sh_link : ehdr->e_shstrndx;
  char *shstrtab = buf + shdrs[shstrndx].sh_offset;
 
  for (int i = 0; i < shnum; i++) {
	Elf64_Shdr *shdr = &shdrs[i];	
 
	if (shdr->sh_type == SHT_DYNSYM) {
		Elf64_Sym *syms = (Elf64_Sym *)(buf + shdr->sh_offset); 
		int entries = shdr->sh_size / shdr->sh_entsize;
		char *strtab = buf + shdrs[shdr->sh_link].sh_offset;

		for (int j = 0; j < entries; j++) {
			Elf64_Sym *sym = &syms[j];
			char *name = strtab + sym->st_name;
			if(name[0] == 0x00) continue;
			if (strlen(name) > sizeof(g_sotable[g_pos].name)) {
			   printf("[load_sym]: funcPtr symbol: %s too big\n", name);
			}
			strcpy(g_sotable[g_pos].name, name);				
			if (sym->st_value) {
				g_sotable[g_pos++].funcPtr = (void *)get_addr(secinfo, sym->st_value);
				//printf("[load_sym]: --0--funcPtr symbol: %s,%p\n", name, sym->st_value);
			} else {
				g_sotable[g_pos++].funcPtr = find_dlso(name);
				//printf("[load_sym]: --1--funcPtr symbol: %s,%p,%p\n", name, sym->st_value,sym->st_size);
			}
		}
	}
  }
  return 0;
}

int do_relocate(char *title, ElfSecInfo *secinfo)
{
  relocate(title, "rel", secinfo);
  relocate(title, "pltrel", secinfo);
  return 0;
}

int load_so(char *fname)
{
  printf("[load_so]: %s\n",fname);

  FILE *fp = fopen(fname, "r+b");
  if (fp == NULL) {
  	printf("[load_so]: not found %s\n", fname);
  	return -1;
  }
  fseek(fp, 0, SEEK_END);
  int size = ftell(fp);
  char *buf = (char *)malloc(size);
  fseek(fp, 0L, SEEK_SET);
  fread(buf, size, 1, fp);
  fclose(fp);

  Elf64_Ehdr *ehdr = (Elf64_Ehdr *)buf;

  if (memcmp(ehdr->e_ident, ELFMAG, SELFMAG)) {
    printf("[load_so]: not elf");
    return -1;
  }

  if (ehdr->e_type != ET_DYN) {
    printf("[load_so]: not so\n");
	free(buf);
    return -1;
  }

  ElfSecInfo  secinfo = {0};
  if (LoadAndSecInfo(buf, ehdr, &secinfo) < 0) {
	return -1;
  }

  load_sym(buf, ehdr, &secinfo);
  do_relocate("load_so", &secinfo);

  free(buf);

  if (secinfo.initEntryAddr) {
	entryFunc func = (entryFunc)get_addr(&secinfo, secinfo.initEntryAddr);
	printf("[load_so1]: init point: %p\n", func);
	func();
  }
	
  if (secinfo.initArrayAddr) {
	for (int i = 0; i < secinfo.initArraySize/sizeof(uint64_t); i++) {
	    uint64_t *addr = (uint64_t *)(secinfo.initArrayAddr+i*sizeof(uint64_t));
		entryFunc func = (entryFunc)get_addr(&secinfo, *addr);
		printf("[load_so1]: initarray point: %p\n", func);
		func();
	}
  }

  entryFunc entryPointFunc = (entryFunc)secinfo.entryAddr;
 
  printf("[load_so]: entry point: %p\n", entryPointFunc);
  printf("====== finish loading so ======\n");

  entryPointFunc();

  return 0;
}

int load_library(ElfSecInfo *secinfo)
{
  char *dynamicSymbolNames = (char *)secinfo->strTabAddr;
  Elf64_Sym* dynamicSymbols = (Elf64_Sym *)secinfo->symTabAddr;

  for (uint32_t i = 0; i < secinfo->soindex; i++) {
	  int find = 0;
	  char *symbolName = dynamicSymbolNames + secinfo->soaddr[i];	
	  //printf("[load_library]: %s\n", symbolName);
	  for (int j = 0; j < sizeof(g_sohandle)/sizeof(so_handle_t); j++) {
		  if (strcmp(symbolName, g_sohandle[j].name) == 0) {
			 find = 1;
			 break;
		  }
	  }
	  if (find == 0) {
		 load_so(symbolName);
	  }		
  }
  return 0;
}

int load_exe(char *fname)
{
  printf("[load_exe]: %s\n", fname);

  FILE *fp = fopen(fname, "r+b");
  if (fp == NULL) {
  	printf("[load_exe]: not found %s\n", fname);
  	return -1;
  }
  fseek(fp, 0, SEEK_END);
  int size = ftell(fp);
  char *buf = (char *)malloc(size);
  fseek(fp, 0L, SEEK_SET);
  fread(buf, size, 1, fp);
  fclose(fp);
 
  Elf64_Ehdr *ehdr = (Elf64_Ehdr *)buf;

  if (memcmp(ehdr->e_ident, ELFMAG, SELFMAG)) {
    printf("[load_exe]: not elf");
	free(buf);
    return -1;
  }

  if (ehdr->e_type != ET_EXEC && ehdr->e_type != ET_DYN) {
    printf("[load_exe]: not exec[%d]\n", ehdr->e_type);
	free(buf);
    return -1;
  }

  ElfSecInfo  secinfo = {0};
  if (LoadAndSecInfo(buf, ehdr, &secinfo) < 0) {
	return -1;
  }

  load_library(&secinfo);

  load_sym(buf, ehdr, &secinfo);
  do_relocate("load_exe", &secinfo);

  free(buf);

  if (secinfo.initEntryAddr) {
	entryFunc func = (entryFunc)get_addr(&secinfo, secinfo.initEntryAddr);
	printf("[load_exe]: init point: %p\n", func);
	func();
  }

  if (secinfo.initArrayAddr) {
	for (int i = 0; i < secinfo.initArraySize/sizeof(uint64_t); i++) {
	    uint64_t *addr = (uint64_t *)(secinfo.initArrayAddr+i*sizeof(uint64_t));
		entryFunc func = (entryFunc)get_addr(&secinfo, *addr);
		printf("[load_exe]: initarray point: %p\n", func);
		func();
	}
  }

  entryFunc entryPointFunc = (entryFunc)get_addr(&secinfo, secinfo.entryAddr);

  printf("[load_exe]: entry point: %p\n", entryPointFunc);
  printf("====== finish loading exe ======\n");

  entryPointFunc();
 
  return 0;
}

int main(int argc, char **argv)
{
  if (load_dl() < 0) {
  	exit(-1);
  }

  g_ftable[0].funcPtr = stdin;
  g_ftable[1].funcPtr = stdout;
  g_ftable[2].funcPtr = stderr;
  g_ftable[3].funcPtr = &sre_optarg;

  g_argc = argc - 1;
  g_argv = argv + 1;

  load_exe(argv[1]);
  
  close_dl();

  return 1;
}
